import cv2
from ultralytics import YOLO
from utils import img_utils

# Load a model
model = YOLO(r"best.pt")  # load a pretrained model (recommended for training)
# Train the model
# cv2读取视频
cap = cv2.VideoCapture("C:/Users/kang_/Desktop/0227大猪/1.mp4")
if cap.isOpened():
    while True:
        ret, frame = cap.read()
        # 如果读取失败或到达视频末尾，退出循环
        if not ret:
            break
        results = model(frame, imgsz=640, conf=0.5)
        for result in results:
            cv2.imshow('result', result.plot())
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
